package main

/*
#cgo windows LDFLAGS: -L./OpenBLAS-0.3.30-x64/lib -l:libopenblas.a -static -lpthread
#cgo windows CFLAGS: -I./OpenBLAS-0.3.30-x64/include
#include <cblas.h>
*/
import "C"
import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"
	"unicode"
)

// 在init函数中设置线程亲和性
func init() {
	runtime.LockOSThread()
	os.Setenv("OPENBLAS_NUM_THREADS", fmt.Sprint(runtime.NumCPU()))
	runtime.GOMAXPROCS(runtime.NumCPU())
}

// 打印 OpenBLAS 配置信息
func printOpenBLASConfig() {
	config := C.GoString(C.openblas_get_config())
	if config == "" {
		fmt.Println("❌ 未检测到 OpenBLAS 配置")
		return
	}

	fmt.Println("OpenBLAS 配置信息:")
	fmt.Println(config)

	// 检查核心架构
	if C.openblas_get_corename() != nil {
		coreName := C.GoString(C.openblas_get_corename())
		fmt.Printf("使用的核心架构: %s\n", coreName)
	}
}

// 通用语料加载函数
func loadCorpus(filePath string, startLine int, segmentText bool, perLine bool, maxSize int) []string {
	file, err := os.Open(filePath)
	if err != nil {
		fmt.Printf("加载文件失败: %v\n", err)
		return []string{}
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	corpus := []string{}
	lineNum := 0

	for scanner.Scan() {
		lineNum++
		// 跳过指定行数之前的内容
		if lineNum < startLine {
			continue
		}

		line := scanner.Text()
		// 去除前后空白
		line = strings.TrimSpace(line)
		if len(line) == 0 {
			continue
		}

		var processedLine string
		if segmentText && !perLine {
			// 处理中英文混合文本，按字分割
			runes := []rune(line)
			var segmented string
			for i := 0; i < len(runes); {
				// 跳过空格
				if runes[i] == ' ' {
					segmented += " "
					i++
					continue
				}

				// 检查是否是中文字符
				if unicode.Is(unicode.Han, runes[i]) {
					// 中文字符按字分割
					segmented += string(runes[i]) + " "
					i++
				} else {
					// 非中文字符按连续字符分割
					j := i
					for j < len(runes) && runes[j] != ' ' && !unicode.Is(unicode.Han, runes[j]) {
						j++
					}
					if j > i {
						segmented += string(runes[i:j]) + " "
					}
					i = j
				}
			}
			// 去除末尾空格
			processedLine = strings.TrimSpace(segmented)
		} else if perLine {
			// 直接使用整行
			processedLine = line
		} else {
			// 不分割文本，直接使用
			processedLine = line
		}

		corpus = append(corpus, processedLine)

		// 限制语料库大小
		if maxSize > 0 && len(corpus) >= maxSize {
			break
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Printf("读取文件错误: %v\n", err)
	}

	return corpus
}

// 从文件加载小说语料，支持从指定行开始读取
func loadNovelCorpus(filePath string, startLine int) []string {
	// 调用通用语料加载函数，设置需要分割文本，最大1000行
	return loadCorpus(filePath, startLine, true, false, 1000)
}

// 从文件加载词典语料，支持从指定行开始读取
func loadDictionary(filePath string, startLine int) []string {
	// 调用通用语料加载函数，设置不需要分割文本，加载所有行
	return loadCorpus(filePath, startLine, false, false, 0)
}

// 训练标记结构
type TrainMark struct {
	Epoch          int     `json:"epoch"`
	Loss           float32 `json:"loss"`
	LR             float32 `json:"lr"`
	CorpusFilePath string  `json:"corpus_file_path"`
	LineNumber     int     `json:"line_number"`
}

// 保存训练标记到文件
func saveTrainMark(mark TrainMark, filePath string) error {
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	encoder := json.NewEncoder(file)
	return encoder.Encode(mark)
}

// 从文件加载训练标记
func loadTrainMark(filePath string) (*TrainMark, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var mark TrainMark
	decoder := json.NewDecoder(file)
	err = decoder.Decode(&mark)
	if err != nil {
		return nil, err
	}

	return &mark, nil
}

func main() {
	printOpenBLASConfig()

	// 定义文件路径
	var modelPath string = "transformer_model.json"
	var trainMarkPath string = "train_mark.json"

	// 初始化训练标记变量
	var trainMark *TrainMark
	var trainMarkErr error = errors.New("no train mark")

	// 确定语料文件路径和起始行数
	corpusFilePath := "./data/data.txt"
	startLine := 1

	// 尝试加载训练标记
	trainMark, trainMarkErr = loadTrainMark(trainMarkPath)

	// 如果有训练标记，使用上次训练的位置
	if trainMarkErr == nil {
		// 检查语料文件路径是否为空
		if trainMark.CorpusFilePath != "" {
			corpusFilePath = trainMark.CorpusFilePath
		}
		startLine = trainMark.LineNumber + 1 // 从下一行开始
		fmt.Printf("将从 %s 的第 %d 行继续训练\n", corpusFilePath, startLine)
	} else {
		fmt.Printf("将从 %s 的第 %d 行开始训练\n", corpusFilePath, startLine)
	}

	// 从文件加载词典语料
	dictionaryFilePath := "./data/isChinese.dictionary.txt"
	dictionaryStartLine := 1
	dictionary := loadDictionary(dictionaryFilePath, dictionaryStartLine)
	fmt.Printf("成功加载词典，包含 %d 个词汇\n", len(dictionary))

	var err error
	// 为词典构建词汇表（不分割中文词组）
	dictVocab, _, err := BuildVocab(dictionary, 1, 2, "id_mapping.json", false)
	if err != nil {
		fmt.Printf("构建词典词汇表失败: %v\n", err)
		return
	}
	fmt.Printf("成功构建词典词汇表，包含 %d 个词汇\n", len(dictVocab))

	// 从文件加载小说语料
	corpus := loadNovelCorpus(corpusFilePath, startLine)
	fmt.Printf("成功加载 %d 条训练语句\n", len(corpus))

	// 为小说语料构建词汇表（按字分割中文），此时会从id_mapping.json读取已有的词典词汇
	vocab, reverseVocab, err := BuildVocab(corpus, 1, 3, "id_mapping.json", true)
	if err != nil {
		fmt.Printf("构建小说词汇表失败: %v\n", err)
		return
	}
	fmt.Printf("成功构建包含词典的词汇表，包含 %d 个词汇\n", len(vocab))

	// 训练并使用Transformer语言模型
	fmt.Println("\n\033[1;34m===== Transformer模型训练与生成 ======\033[0m")

	var loadedModel *TransformerLanguageModel

	// 尝试加载训练标记
	trainMark, trainMarkErr = loadTrainMark(trainMarkPath)
	needTrain := true

	// 尝试加载模型
	if _, err = os.Stat(modelPath); err == nil {
		fmt.Printf("\033[1;36m正在从 %s 加载模型...\033[0m\n", modelPath)
		loadStart := time.Now()
		loadedModel, err = LoadTransformerModel(modelPath)
		if err != nil {
			fmt.Printf("\033[1;31m加载模型失败: %v\033[0m\n", err)
		} else {
			fmt.Printf("\033[1;32m模型加载成功!\033[0m 耗时: %s\n", time.Since(loadStart))

			// 检查是否需要训练
			if trainMarkErr == nil && trainMark.Loss <= 0.005 {
				fmt.Printf("\033[1;32m检测到训练标记 (Loss: %.4f)，不需要重新训练\033[0m\n", trainMark.Loss)
				needTrain = false
			} else {
				fmt.Printf("\033[1;33m训练标记不存在或Loss值过高，需要继续训练\033[0m\n")
			}
		}
	} else {
		fmt.Printf("\033[1;33m模型文件不存在，需要训练\033[0m\n")
	}

	// 如果需要训练
	if needTrain {
		var transformerModel *TransformerLanguageModel

		// 尝试加载已有的模型（如果存在）
		modelFileInfo, err := os.Stat(modelPath)
		var trainMark *TrainMark
		var trainMarkErr error
		startEpoch := 0
		learningRate := float32(0.03)

		// 尝试加载模型
		if err == nil && modelFileInfo.Size() > 0 {
			fmt.Printf("\033[1;36m检测到模型文件 %s (大小: %d 字节)，尝试加载以继续训练...\033[0m\n", modelPath, modelFileInfo.Size())
			loadStart := time.Now()
			transformerModel, err = LoadTransformerModel(modelPath)
			if err != nil {
				fmt.Printf("\033[1;31m加载模型失败: %v，将创建新模型\033[0m\n", err)
			} else {
				fmt.Printf("\033[1;32m模型加载成功，可以继续训练!\033[0m 耗时: %s\n", time.Since(loadStart))
				// 重要：更新模型的词汇表为当前构建的词汇表
				transformerModel.Vocab = vocab
				transformerModel.ReverseVocab = reverseVocab
				transformerModel.VocabSize = len(vocab)

				// 加载训练标记以获取上次训练的epoch数
				trainMark, trainMarkErr = loadTrainMark(trainMarkPath)
				if trainMarkErr == nil {
					fmt.Printf("\033[1;32m检测到上次训练标记: 完成 %d 个epoch, Loss: %.4f\033[0m\n", trainMark.Epoch, trainMark.Loss)
				} else {
					fmt.Printf("\033[1;33m未找到训练标记，将从头开始训练\033[0m\n")
				}
			}
		} else {
			if err != nil {
				fmt.Printf("\033[1;33m模型文件不存在: %v，将创建新模型\033[0m\n", err)
			} else {
				fmt.Printf("\033[1;33m模型文件为空，将创建新模型\033[0m\n")
			}
		}

		// 创建新模型（如果加载失败或模型不存在）
		if transformerModel == nil {
			transformerModel = NewTransformerLanguageModelWithVocab(vocab, reverseVocab)
			if transformerModel == nil {
				fmt.Printf("\033[1;31m创建新模型失败！程序无法继续\033[0m\n")
				return
			} else {
				fmt.Printf("\033[1;32m新模型创建成功\033[0m\n")
			}
		}

		// 训练模型
		fmt.Printf("\033[1;36m开始训练模型...\033[0m\n")
		totalEpochs := 15 // 设置总训练轮次为15

		// 如果有上次训练的标记，使用上次的学习率和epoch数
		if trainMarkErr == nil && trainMark != nil {
			learningRate = trainMark.LR
			startEpoch = trainMark.Epoch
			fmt.Printf("\033[1;32m从上次训练的 epoch %d, 学习率 %.6f 继续训练\033[0m\n", startEpoch, learningRate)
		} else if trainMarkErr == nil {
			fmt.Printf("\033[1;33m训练标记存在但内容无效，将从头开始训练\033[0m\n")
		}

		fmt.Println("\033[1;36m===== 使用小说语料训练 =====\033[0m")
		fmt.Printf("\033[1;34m使用小说语料进行 %d 个epoch的训练...\033[0m\n", totalEpochs-startEpoch)
		fmt.Printf("\033[1;32m使用初始学习率 %.6f 开始训练\033[0m\n", learningRate)

		startTime := time.Now()
		lastLoss, lastLR, _ := transformerModel.Train(corpus, totalEpochs, learningRate, startEpoch, 5, 0.001)
		trainDuration := time.Since(startTime)
		fmt.Printf("\033[1;32m模型训练完成!\033[0m 耗时: %s\n", trainDuration)

		// 确定最后处理的行数
		lastLineNumber := startLine + len(corpus) - 1

		// 保存训练标记
		trainMarkValue := TrainMark{
			Epoch:          totalEpochs,
			Loss:           lastLoss,
			LR:             lastLR,
			CorpusFilePath: corpusFilePath,
			LineNumber:     lastLineNumber,
		}
		fmt.Printf("\033[1;36m正在保存训练标记到 %s...\033[0m\n", trainMarkPath)
		err = saveTrainMark(trainMarkValue, trainMarkPath)
		if err != nil {
			fmt.Printf("\033[1;31m保存训练标记失败: %v\033[0m\n", err)
		} else {
			fmt.Printf("\033[1;32m训练标记保存成功!\033[0m\n")
		}

		// 保存模型
		fmt.Printf("\033[1;36m正在保存模型到 %s...\033[0m\n", modelPath)
		saveStart := time.Now()
		err = transformerModel.SaveTransformerModel(modelPath)
		if err != nil {
			fmt.Printf("\033[1;31m保存模型失败: %v\033[0m\n", err)
			return
		}
		fmt.Printf("\033[1;32m模型保存成功!\033[0m 耗时: %s\n", time.Since(saveStart))

		// 模型评估
		fmt.Println("\033[1;36m===== 模型评估 =====\033[0m")
		// 从语料中分割出10%作为验证集
		valSize := len(corpus) / 10
		if valSize == 0 && len(corpus) > 0 {
			valSize = 1
		}
		valCorpus := corpus[len(corpus)-valSize:]
		// 评估模型
		valLoss := loadedModel.Evaluate(valCorpus)
		fmt.Printf("\033[1;32m验证集Loss: %.4f\033[0m\n", valLoss)

		// 加载刚保存的模型
		fmt.Printf("\033[1;36m正在从 %s 加载模型...\033[0m\n", modelPath)
		loadStart := time.Now()
		loadedModel, err = LoadTransformerModel(modelPath)
		if err != nil {
			fmt.Printf("\033[1;31m加载模型失败: %v\033[0m\n", err)
			return
		}
		fmt.Printf("\033[1;32m模型加载成功!\033[0m 耗时: %s\n", time.Since(loadStart))

	}

	// 使用加载的模型生成文本
	prompt := "赵祯在御座上" // 使用词汇表中存在的字词作为提示
	fmt.Println("\n[信息] 使用加载的模型生成文本:")

	// 使用累积的文本作为新的提示词进行最终生成
	finalGenStart := time.Now()
	finalGeneratedText, _ := loadedModel.Thinking(prompt, 100, false)
	finalGenDuration := time.Since(finalGenStart)

	// 统计最终生成结果中的unk数量
	finalUnkCount := 0
	finalWords := strings.Fields(finalGeneratedText)
	for _, word := range finalWords {
		if word == "<unk>" {
			finalUnkCount++
		}
	}

	// 格式化输出最终结果，使用Windows命令行兼容的颜色代码
	fmt.Printf("[最终结果] %s\n", finalGeneratedText)
	fmt.Printf("  [统计] 生成时间: %s | 未知/重复标记数量: %d (%d%%)\n", finalGenDuration, finalUnkCount, (finalUnkCount*100)/len(finalWords))

	// 进入循环输入模式
	fmt.Println("\n[提示] 进入交互模式，输入内容后按回车生成文本，输入 'exit' 退出...")
	for {
		fmt.Print("\n请输入: ")
		var input string
		fmt.Scanln(&input)

		if input == "exit" {
			fmt.Println("退出程序...")
			break
		}

		if len(input) == 0 {
			fmt.Println("输入不能为空，请重新输入。")
			continue
		}

		// 调用Thinking方法生成文本
		genStart := time.Now()
		answer, thoughtProcess := loadedModel.Thinking(input, 200, true)
		genDuration := time.Since(genStart)

		// 统计unk数量
		unkCount := 0
		words := strings.Fields(answer)
		for _, word := range words {
			if word == "<unk>" {
				unkCount++
			}
		}

		// 格式化输出
		fmt.Printf("\n[思考过程] %s\n", thoughtProcess)
		fmt.Printf("\n[生成结果] %s\n", answer)
		fmt.Printf("  [统计] 生成时间: %s | 未知/重复标记数量: %d (%d%%)\n", genDuration, unkCount, (unkCount*100)/len(words))
		fmt.Println("--------------------------------------------------")
	}
}
